# app.py
import os, requests, json
import datetime as dt
import tempfile
from flask import Flask, render_template, request, redirect, url_for, flash, jsonify, Response
from flask_sqlalchemy import SQLAlchemy
from flask_login import LoginManager, login_user, logout_user, login_required, current_user, UserMixin
from werkzeug.security import generate_password_hash, check_password_hash
from openai import OpenAI
from config import API_CONFIG, SECRET_KEY, DEEPSEEK_MODEL
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader, PyPDFLoader, UnstructuredWordDocumentLoader, UnstructuredPowerPointLoader, CSVLoader
from langchain_community.tools import DuckDuckGoSearchRun

# 设置环境变量
os.environ["HF_HUB_OFFLINE"] = "1"  # 禁止任何 HTTP 请求
os.environ["TRANSFORMERS_OFFLINE"] = "1"  # 让 transformers 也不联网
os.environ["HF_HOME"] = "./model"  # 把缓存目录指到项目内
os.environ['TZ'] = 'Asia/Shanghai'  # 把系统时区强制设为上海

# 打印当前工作目录和临时目录
print("cwd:", os.getcwd())
print("temp:", tempfile.gettempdir())

# 初始化 Flask 应用
app = Flask(__name__)
app.config['SECRET_KEY'] = SECRET_KEY
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///' + os.path.join(
    os.path.abspath(os.path.dirname(__file__)), 'chat.db'
)
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

db = SQLAlchemy(app)
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = 'login'

client = OpenAI(api_key=API_CONFIG["api_key"], base_url=API_CONFIG["base_url"])
MODEL = DEEPSEEK_MODEL
print(">>> 当前 LLM 模型:", MODEL)

# 数据模型
class User(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    password_hash = db.Column(db.String(128), nullable=False)
    is_admin = db.Column(db.Boolean, default=False)

    def set_password(self, pwd):
        self.password_hash = generate_password_hash(pwd)

    def check_password(self, pwd):
        return check_password_hash(self.password_hash, pwd)

class Message(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    body = db.Column(db.Text, nullable=False)
    author_id = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=False)
    author = db.relationship('User', backref='messages')
    created = db.Column(
        db.DateTime,
        default=lambda: dt.datetime.now(dt.timezone(dt.timedelta(hours=8)))
    )
    ai_reply = db.Column(db.Text)

# 登录回调
@login_manager.user_loader
def load_user(user_id):
    return db.session.get(User, int(user_id))

# 初始化数据库
with app.app_context():
    db.create_all()
    if not User.query.filter_by(username='admin').first():
        admin = User(username='admin', is_admin=True)
        admin.set_password('admin')
        db.session.add(admin)
        db.session.commit()

# 注册表
LOADER_MAP = {
    ".txt": TextLoader,
    ".md": TextLoader,
    ".csv": CSVLoader,
    ".pdf": PyPDFLoader,
    ".docx": UnstructuredWordDocumentLoader,
    ".pptx": UnstructuredPowerPointLoader,
    ".html": TextLoader,
    ".htm": TextLoader,
}

# 自定义加载器选择函数
def load_document(file_path: str):
    ext = os.path.splitext(file_path)[1].lower()
    loader_class = LOADER_MAP.get(ext)
    if not loader_class:
        return None

    try:
        if loader_class == TextLoader:
            loader = loader_class(file_path, encoding="utf-8")
        else:
            loader = loader_class(file_path)
        return loader.load()
    except Exception as e:
        print(f"加载文件 {file_path} 时出错: {e}")
        return None

# 修改后的向量库构建函数
def build_or_load_vector_store():
    persist_dir = "knowledge_db"
    embeddings = HuggingFaceEmbeddings(model_name="./model/all-MiniLM-L6-v2")

    if os.path.exists(persist_dir):
        print("加载现有的向量数据库...")
        try:
            return Chroma(persist_directory=persist_dir, embedding_function=embeddings)
        except Exception as e:
            print(f"加载现有向量库失败: {e}，重新构建...")
            import shutil
            shutil.rmtree(persist_dir)

    print("构建新的向量数据库...")
    docs = []
    my_docs_dir = "my_docs"

    if not os.path.exists(my_docs_dir):
        print(f"警告: {my_docs_dir} 目录不存在，创建空目录")
        os.makedirs(my_docs_dir, exist_ok=True)
        vectorstore = Chroma.from_documents([], embeddings, persist_directory=persist_dir)
        print("空向量数据库创建完成")
        return vectorstore

    for root, _, files in os.walk(my_docs_dir):
        for file in files:
            file_path = os.path.join(root, file)
            file_ext = os.path.splitext(file)[1].lower()

            if file_ext in LOADER_MAP:
                print(f"正在加载: {file_path}")
                loaded_docs = load_document(file_path)
                if loaded_docs:
                    docs.extend(loaded_docs)
            else:
                print(f"跳过不支持的文件格式: {file_path}")

    if not docs:
        print("未找到可处理的文档，创建空向量库")
        vectorstore = Chroma.from_documents([], embeddings, persist_directory=persist_dir)
        print("空向量数据库创建完成")
        return vectorstore

    splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
    chunks = splitter.split_documents(docs)

    print(f"共处理 {len(docs)} 个文档，分割为 {len(chunks)} 个片段")
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=persist_dir
    )
    print("向量数据库构建完成并已自动持久化")
    return vectorstore

try:
    vector_store = build_or_load_vector_store()
    print("向量库加载成功")
except Exception as e:
    print(f"向量库加载失败: {e}")
    MODEL_DIR = os.getenv("MODEL_PATH") or os.path.join(
        os.path.abspath(os.path.dirname(__file__)),
        "model",
        "all-MiniLM-L6-v2"
    )
    embeddings = HuggingFaceEmbeddings(model_name=MODEL_DIR, cache_folder=None)
    vector_store = None
    print("使用空向量库作为后备")

# 搜索函数
ddg = DuckDuckGoSearchRun()

def search(query_list, num=6):
    api_key = os.getenv("TAVILY_API_KEY")
    if not api_key:
        print("警告：未设置 TAVILY_API_KEY，直接回退到 DuckDuckGo")
        return ddg_fallback(" ".join(query_list), num)

    url = "https://api.tavily.com/search"
    payload = {
        "query": " ".join(query_list),
        "max_results": num,
        "search_depth": "basic",
        "include_answer": False,
        "include_raw_content": True
    }
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    try:
        r = requests.post(url, json=payload, headers=headers, timeout=8)
        r.raise_for_status()
        return r.json()
    except Exception as e:
        print("Tavily 搜索异常:", e)
        return ddg_fallback(" ".join(query_list), num)

def ddg_fallback(query, num):
    try:
        results = ddg.run(query)
        snippets = [r.strip() for r in results.split("\n") if r.strip()][:num]
        return {
            "results": [
                {"title": f"结果{i + 1}", "url": "", "content": s, "snippet": s}
                for i, s in enumerate(snippets)
            ]
        }
    except Exception as e:
        print("DuckDuckGo 也挂了:", e)
        return {"error": str(e), "results": []}

def format_search_results(results):
    if not results or "error" in results:
        return "未能获取实时搜索结果"

    if "results" not in results or not results["results"]:
        return "未找到相关搜索结果"

    formatted = []
    for item in results["results"][:4]:
        title = item.get("title", "无标题")
        url = item.get("url", "")
        content = item.get("content", "") or item.get("snippet", "")
        formatted.append(f"标题：{title}\n链接：{url}\n摘要：{content[:500]}…")
    return "\n\n".join(formatted)

# 路由：认证
@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form['username'].strip()
        pwd = request.form['password']
        if User.query.filter_by(username=username).first():
            flash('用户名已存在')
            return redirect(url_for('register'))
        u = User(username=username)
        u.set_password(pwd)
        db.session.add(u)
        db.session.commit()
        flash('注册成功，请登录')
        return redirect(url_for('login'))
    return render_template('register.html')

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        u = User.query.filter_by(username=request.form['username']).first()
        if u and u.check_password(request.form['password']):
            login_user(u)
            return redirect(url_for('index'))
        flash('用户名或密码错误')
    return render_template('login.html')

@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('login'))

# 路由：留言板
@app.route('/')
@login_required
def index():
    msgs = Message.query.order_by(Message.created.desc()).all()
    return render_template('index.html', msgs=msgs)

@app.route('/add', methods=['POST'])
@login_required
def add_message():
    body = request.form['body'].strip()
    if body:
        m = Message(body=body, author=current_user)
        db.session.add(m)
        db.session.commit()
    return redirect(url_for('index'))

# 路由：管理员
@app.route('/reply/<int:msg_id>')
def reply_stream(msg_id):
    msg = Message.query.get_or_404(msg_id)

    if msg.ai_reply:
        return jsonify({'content': msg.ai_reply})

    MODEL_DIR = os.getenv("MODEL_PATH") or os.path.join(
        os.path.abspath(os.path.dirname(__file__)),
        "model",
        "all-MiniLM-L6-v2"
    )
    embeddings = HuggingFaceEmbeddings(model_name=MODEL_DIR, cache_folder=None)

    if vector_store:
        current_vector_store = vector_store
    else:
        try:
            current_vector_store = Chroma(
                embedding_function=embeddings,
                persist_directory="knowledge_db",
                collection_metadata={"hnsw:space": "cosine"}
            )
        except Exception as e:
            print(f"创建空向量库失败: {e}")
            return jsonify({'content': '知识库初始化失败，请联系管理员'})

    def generate():
        full = ''
        try:
            if current_vector_store:
                docs = current_vector_store.similarity_search(msg.body, k=20)
                private_ctx = "\n".join([d.page_content for d in docs])
            else:
                private_ctx = "知识库暂不可用"

            web_ctx = format_search_results(search([msg.body]))

            user_prompt = f"【公司知识】{private_ctx}\n【实时搜索】{web_ctx}\n用户问题：{msg.body}"

            stream = client.chat.completions.create(
                model=MODEL,
                messages=[
                    {"role": "system", "content": _make_system_prompt()},
                    {"role": "user", "content": user_prompt}
                ],
                stream=True,
                temperature=0.5,
                max_tokens=2048,
                top_p=0.95
            )

            for chunk in stream:
                delta = chunk.choices[0].delta.content or ""
                if delta:
                    full += delta
                    yield f"data: {json.dumps({'content': delta})}\n\n"

            with app.app_context():
                msg.ai_reply = full
                db.session.commit()
        except Exception as e:
            error_msg = f"生成回复时出错: {str(e)}"
            print(error_msg)
            yield f"data: {json.dumps({'content': error_msg, 'error': True})}\n\n"

    return Response(generate(), mimetype='text/event-stream')

@app.route('/admin')
@login_required
def admin():
    if not current_user.is_admin:
        flash('无权限')
        return redirect(url_for('index'))
    users = User.query.all()
    return render_template('admin.html', users=users)

@app.route('/admin/toggle/<int:user_id>')
@login_required
def toggle_admin(user_id):
    if not current_user.is_admin:
        return 'Forbidden', 403
    u = User.query.get_or_404(user_id)
    u.is_admin = not u.is_admin
    db.session.commit()
    return redirect(url_for('admin'))

# 路由：文件上传
@app.route('/upload', methods=['POST'])
@login_required
def upload_file():
    from werkzeug.utils import secure_filename
    file = request.files.get('file')
    if not file or file.filename == '':
        return jsonify({'status': 'error', 'msg': '未选择文件'}), 400
    filename = secure_filename(file.filename)
    ext = os.path.splitext(filename)[1].lower()
    if ext not in LOADER_MAP:
        return jsonify({'status': 'error', 'msg': f'不支持的文件类型: {ext}'}), 400
    os.makedirs("my_docs", exist_ok=True)
    save_path = os.path.join("my_docs", filename)
    file.save(save_path)
    try:
        global vector_store
        vector_store = build_or_load_vector_store()
    except Exception as e:
        return jsonify({'status': 'error', 'msg': f'知识库更新失败: {str(e)}'}), 500
    return jsonify({'status': 'ok', 'msg': '上传并更新成功'})

# 工具函数
def _make_system_prompt():
    today = dt.datetime.now(dt.timezone(dt.timedelta(hours=8))).date()
    weekday = ['一', '二', '三', '四', '五', '六', '日'][today.weekday()]
    return f"""
You are DeepSeek Chat, a large language model trained by DeepSeek, based on the DeepSeek-V3.1 architecture.
Knowledge cutoff: 2024-07
今天是 {today.year}年{today.month}月{today.day}日，星期{weekday}。请始终用这句话回答任何关于当前日期的问题。
You can use tools like web search, drawing, code interpreter, etc., when needed.
Answer in the same language as the user.
You are a helpful assistant. Always answer accurately and concisely.
"""

def check_configurations():
    required_env_vars = ["SECRET_KEY"]
    missing_vars = [var for var in required_env_vars if not os.getenv(var)]
    if missing_vars:
        print(f"警告: 缺少环境变量: {', '.join(missing_vars)}")
    if not os.path.exists("my_docs"):
        print("警告: my_docs 目录不存在，知识库功能将不可用")
    if not os.path.exists("knowledge_db"):
        print("信息: 知识向量库不存在，将在首次请求时创建")

if __name__ == '__main__':
    print('进入 __main__')
    check_configurations()
    app.run(host='0.0.0.0', port=5000, debug=False)